{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Multimodal (Vision) Finetuning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import base64\n",
    "import os\n",
    "from os import PathLike\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tensorzero import AsyncTensorZeroGateway\n",
    "from tqdm.asyncio import tqdm_asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = Path(\"data\")\n",
    "\n",
    "LABELS_PATH = DATA_PATH / \"labels.csv\"\n",
    "\n",
    "CONCURRENCY = 10\n",
    "\n",
    "VARIANT_NAME = \"baseline\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(path: PathLike):\n",
    "    assert LABELS_PATH.exists(), (\n",
    "        f\"Labels file {LABELS_PATH} does not exist. See the README.md and ensure you've downloaded the dataset correctly.\"\n",
    "    )\n",
    "\n",
    "    df = pd.read_csv(LABELS_PATH)\n",
    "\n",
    "    # Sanity Check: ensure every image exists\n",
    "    for _, row in df.iterrows():\n",
    "        img_path = path / Path(row[\"document\"])\n",
    "        assert img_path.exists(), (\n",
    "            f\"Image {img_path} does not exist. See the README.md and ensure you've downloaded the dataset correctly.\"\n",
    "        )\n",
    "\n",
    "    train_df = df[df[\"is_train\"] == 1].reset_index(drop=True)\n",
    "    test_df = df[df[\"is_train\"] == 0].reset_index(drop=True)\n",
    "\n",
    "    return train_df, test_df\n",
    "\n",
    "\n",
    "train_df, test_df = load_data(DATA_PATH)\n",
    "\n",
    "print(f\"Found {len(train_df)} train documents and {len(test_df)} test documents\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"tensorzero/object_storage\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = await AsyncTensorZeroGateway.build_http(\n",
    "    gateway_url=\"http://localhost:3000\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_document(path: PathLike):\n",
    "    \"\"\"Load an image and encode as a base64 string\"\"\"\n",
    "    path = DATA_PATH / path\n",
    "    assert path.exists()\n",
    "    assert path.suffix.lower() == \".png\"\n",
    "\n",
    "    with open(path, \"rb\") as image_file:\n",
    "        return base64.b64encode(image_file.read()).decode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def process_document(row):\n",
    "    response = await t0.inference(\n",
    "        function_name=\"classify_document\",\n",
    "        input={\n",
    "            \"system\": \"Categorize this document.\",\n",
    "            \"messages\": [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": [\n",
    "                        {\n",
    "                            \"type\": \"image\",\n",
    "                            \"mime_type\": \"image/png\",\n",
    "                            \"data\": load_document(row.document),\n",
    "                        },\n",
    "                    ],\n",
    "                }\n",
    "            ],\n",
    "        },\n",
    "        dryrun=not row.is_train,\n",
    "        cache_options={\n",
    "            \"enabled\": \"on\",\n",
    "        },\n",
    "        variant_name=VARIANT_NAME,\n",
    "    )\n",
    "\n",
    "    predicted_category = response.output.parsed[\"category\"]\n",
    "    correct_classification = predicted_category == row.label\n",
    "\n",
    "    if row.is_train:\n",
    "        await t0.feedback(\n",
    "            metric_name=\"correct_classification\",\n",
    "            value=correct_classification,\n",
    "            inference_id=response.inference_id,\n",
    "        )\n",
    "\n",
    "        await t0.feedback(\n",
    "            metric_name=\"demonstration\",\n",
    "            value={\n",
    "                \"category\": row.label,\n",
    "            },\n",
    "            inference_id=response.inference_id,\n",
    "        )\n",
    "\n",
    "    return correct_classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "semaphore = asyncio.Semaphore(CONCURRENCY)\n",
    "\n",
    "\n",
    "async def process_document_with_semaphore(row):\n",
    "    async with semaphore:\n",
    "        return await process_document(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = await tqdm_asyncio.gather(*[process_document_with_semaphore(row) for _, row in train_df.iterrows()])\n",
    "\n",
    "print(f\"Train Set Accuracy: {np.mean(scores):.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = await tqdm_asyncio.gather(*[process_document_with_semaphore(row) for _, row in test_df.iterrows()])\n",
    "\n",
    "print(f\"Test Set Accuracy: {np.mean(scores):.1%}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
